#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <pthread.h>

#include "mytbf.h"

/*
 * 定义结构体类型
 */
struct mytbf_st
{
	int cps;
	int burst;
	int token;
	int pos;
	pthread_mutex_t mutex;
};
typedef void (*sighandler_t)(int);

/*
 * 全局变量
 */
static struct mytbf_st * job[MYTBF_MAX];
static pthread_mutex_t mut_job  = PTHREAD_MUTEX_INITIALIZER;
static sighandler_t alrm_handler_save;
pthread_t tid_alrm;
static pthread_once_t init_once = PTHREAD_ONCE_INIT;
/*
 * 获取空位置小函数
 */
static int get_free_pos_unlocked(void)
{
	int i;
	for(i = 0; i < MYTBF_MAX; i++)
	{
		if(job[i] == NULL)
			return i;
	}
	return -1;
}
static void * thr_alrm(void * p)
{
	int i;
	while(1)
	{
		pthread_mutex_lock(&mut_job);
		for(i = 0; i <MYTBF_MAX; i++)
		{
			if(job[i] != NULL)
			{
				pthread_mutex_lock(&job[i]->mutex);
				job[i]->token += job[i]->cps;
				if(job[i]->token > job[i]->burst)
					job[i]->token = job[i]->burst;
				pthread_mutex_unlock(&job[i]->mutex);
			}
		}
		pthread_mutex_unlock(&mut_job);
		usleep(1000 * 1000);
	}
}
static void module_unload()
{
	pthread_cancel(tid_alrm);
	pthread_join(tid_alrm, NULL);

	int i;
	for(i = 0; i < MYTBF_MAX; i++)
	{
		if(job[i] != NULL)
		{
			mytbf_destroy(job[i]);
		}
	}

	pthread_mutex_destroy(&mut_job);
}
static void module_load()
{
	int err;

	err = pthread_create(&tid_alrm, NULL, thr_alrm, NULL);
	if(err)
	{
		fprintf(stderr, "create:%s\n", strerror(err));
	}
	atexit(module_unload);
}


/*
 * 令牌桶初始化函数
 * 或令牌桶添加函数
 */
mytbf_t * mytbf_init(int cps, int burst)
{
	struct mytbf_st * me;
	int pos;

	pthread_once(&init_once, module_load);

	//申请内存存放令牌桶
	me = malloc(sizeof(*me));
	if(me == NULL)
		return NULL;
	//初始化令牌桶的值
	me->token = 0;
	me->cps = cps;
	me->burst = burst;
	pthread_mutex_init(&me->mutex, NULL);

	pthread_mutex_lock(&mut_job);
	//获取数组中的空位置
	pos = get_free_pos_unlocked();
	if(pos < 0)
	{
		pthread_mutex_unlock(&mut_job);
		free(me);
		return NULL;
	}
	//初始化令牌桶的值
	me->pos = pos;
	//把令牌桶的地址放入数组
	job[pos] = me;
	pthread_mutex_unlock(&mut_job);


	//返回令牌桶的地址
	return (mytbf_t *)me;
}
/*
 * 获取较小的数函数
 */
static int min(int a, int b)
{
	if(a < b)
		return a;
	return b;
}
/*
 * 消耗令牌函数
 */
int mytbf_fetchtoken(mytbf_t * ptr, int size)
{
	int n;
	struct mytbf_st * me = ptr;
	//参数检验
	if(size <= 0)
		return -EINVAL;

	pthread_mutex_lock(&me->mutex);
	//等待交付令牌
	while(me->token <= 0)
	{
		pthread_mutex_unlock(&me->mutex);
		sched_yield();
		pthread_mutex_lock(&me->mutex);
	}
	//记录消耗的令牌数
	n = min(me->token, size);
	me->token -= n;
	pthread_mutex_unlock(&me->mutex);

	return n;
}
/*
 * 增长令牌函数
 */
int mytbf_returntoken(mytbf_t * ptr, int size)
{
	struct mytbf_st * me = ptr;
	//参数检验
	if(size <= 0)
		return -EINVAL;
	
	pthread_mutex_lock(&me->mutex);
	//增加令牌
	me->token += size;
	if(me->token > me->burst)
		me->token = me->burst;
	pthread_mutex_unlock(&me->mutex);

	return size;
}
/*
 * 销毁令牌桶
 */
int mytbf_destroy(mytbf_t *ptr)
{
	struct mytbf_st * me = ptr;

	pthread_mutex_lock(&mut_job);
	job[me->pos] = NULL;
	pthread_mutex_unlock(&mut_job);
	
	pthread_mutex_destroy(&me->mutex);
	free(me);
	
	return 0;
}

